前幾天看了 attention 後端的 Python 到 C++ 流程，也解釋 linear、norm、activation、RoPE。今天會看 decode attention 和 extend attention。它們都在做 attention，但資料流、分塊大小、KV cache 讀取方式和 softmax 累積方式不一樣。
req_to_token、seq_lens 怎麼決定 kernel 要讀哪裡LlamaAttention.forward
-> self.attn(q, k, v, forward_batch)
-> RadixAttention.forward
-> forward_batch.attn_backend.forward
-> RVVAttnBackend.forward_decode / forward_extend
-> torch.ops.sgl_kernel.decode_attention_cpu / extend_attention_cpu
-> sgl-kernel/csrc/cpu/riscv64/decode.cpp / extend.cpp
LlamaAttention.forward 會完成 Q/K/V projection、RoPE，再呼叫 self.attn(...)。它不在這裡判斷 RVV。SGLang 啟動 CPU backend 時會先檢查 AMX / RVV 支援；選到 rvv 後，ForwardBatch 會帶著建立好的 RVVAttnBackend。AttentionBackend.forward 再依 forward_mode 呼叫 forward_decode 或 forward_extend。
torch.ops.sgl_kernel.* 已經連到安裝好的 C++ extension。我目前會在啟動推論前編譯 sgl-kernel，環境有 clang-19 時優先使用 clang-19 / clang++-19來編譯 RVV Intrinsic 寫的 AI 運算子，每個 token 的 forward 只會呼叫已載入的 operator，不會重新編譯 C++ kernel。
Decode attention 用在產生下一個 token。每個請求已經有一段 KV cache，這一輪有新的 Q/K/V，要把新 K/V 寫進 cache，再用 Q 加到整段前綴。
簡化流程
new q, new k, new v
-> write new k/v into KV cache
-> read old + new K from KV cache
-> score = q @ K
-> softmax(score)
-> output = softmax(score) @ V
Extend attention 用在 prefill 或 chunked prefill，它一次處理多個新 token。它要同時看 prefix cache 和這一輪 extend 的 token。
簡化流程
q_extend, k_extend, v_extend
prefix K/V in KV cache
-> attention over prefix
-> attention over current extend tokens
-> write output
Decode 比較像一個新 query 去掃既有 KV cache，Extend 比較像一小塊 Q 對 前綴和本輪 K/V 做 tiled attention。
FlashAttention 的核心是 IO-aware tiling。它把 Q、K、V 分塊，在小型 score tile 上更新 online softmax，不把完整的 attention score matrix 寫回主記憶體。
extend.cpp 採用相同的演算法結構，可以稱為 FlashAttention-style tiled attention：
BLOCK_M 決定一次處理幾列 query。BLOCK_N 決定一次讀幾個 key / value token。s_i[BLOCK_M x BLOCK_N] 只保存目前的 score tile。m_acc、l_acc、v_prime 跨 tile 保存 running max、exponential sum 與 value accumulator。Decode 時通常只有一列 query。只切 Q tile 很難提供足夠的平行工作，因此 decode.cpp 會再把長 KV sequence 切成 num_kv_splits。每個 split 獨立掃過自己的 KV block，產生 partial output 和 log-sum-exp，decode_accumulate_kv_splits 最後依 long-sum-exp(LSE) 重新加權合併。這是 FlashDecoding-style split-KV attention 的主要結構。
左圖的 KV splits 可以交給不同 CPU 工作項目平行計算，RVV 負責每個 split 裡的 QK、softmax reduction 與 SV。右圖則在同一個 query tile 上反覆讀取 K/V tile，用 online softmax 更新同一組 accumulator。
我這裡實作了 FlashAttention 與 FlashDecoding 的演算法結構，再用 RVV intrinsic、CPU thread、SGLang 的非連續 KV cache layout 完成硬體對應。
attention kernel 不能只看 tensor 本身。SGLang推論服務裡有請求 pool 和token pool， kernel 要透過 映射 tensor 找到每個請求的 token 位置。
幾個重要 tensor
req_to_token: [max_num_reqs, max_context_len]
req_pool_indices: [num_seqs]
seq_lens: [num_seqs]
loc: [num_seqs]
req_pool_indices[b] 會告訴 kernel 第 b 個 sequence 對應到請求 pool 的哪一列
req_to_token[req_idx, t] 會告訴 kernel 第 t 個 token 在 KV cache 裡的位置
seq_lens[b] 告訴 kernel 這個 sequence 現在有多長
loc[b] 是 decode 這輪新 K/V 要寫入的 cache 位置
所以 C++ kernel 的參數看起來很多
decode_attention_cpu(
query,
k_buffer,
v_buffer,
output,
key,
value,
loc,
attn_logits,
req_to_token,
req_pool_indices,
seq_lens,
sm_scale,
logit_cap)
這些參數都在描述 SGLang 執行環境的請求狀態，RVV kernel 需要靠它們找到每個請求對應的 KV token。
decode.cpp 的入口decode_attention_cpu 會先檢查形狀
CHECK_DIM(3, query);
CHECK_DIM(3, k_buffer);
CHECK_DIM(3, v_buffer);
CHECK_DIM(3, key);
CHECK_DIM(3, value);
它期待
query: [num_tokens, num_heads, head_size]
k_buffer: [max_total_num_tokens, num_heads_kv, head_size]
v_buffer: [max_total_num_tokens, num_heads_kv, head_size_v]
attn_logits: [num_seqs, num_heads, num_kv_splits, head_size_v + 1]
接著先把這輪新的 K/V 寫進 cache
decode_set_kv_buffer(
k_buffer,
v_buffer,
key,
value,
loc,
num_seqs,
num_heads_kv,
head_size,
head_size_v,
...);
然後依照 MHA 或 GQA 路徑呼叫不同 kernel
if (num_heads == num_heads_kv) {
decode_attention_kernel_impl<...>(...);
} else {
decode_attention_grouped_kernel_impl<...>(...);
}
MHA 裡 query head 和 KV head 一樣多。GQA/MQA 裡多個 query head 會共用較少的 KV head，所以 grouped kernel 需要處理 head grouping。
decode kernel 會把 KV sequence 切成 block。每個 block 裡，先透過 req_to_token 拿到這段 token 對應的 KV cache 位置
const index_t* cur_indices =
req_to_token + req_pool_id * max_context_len + n;
接著做
s_i = Q @ K_block
這裡會走 index_gemm_kernel_nt。nt 可以讀成 A normal、B transposed 的 GEMM 形狀。K 來自 KV cache，並且透過 cur_indices 做間接索引。
BF16 / FP16 路徑裡，tiny GEMM 會看到 RVV intrinsic
for (int64_t k = 0; k < K; k += vl) {
vl = __riscv_vsetvl_e32m1(K - k);
...
vc[m * COLS + n] =
__riscv_vfmacc_vv_f32m1_tu(vc[m * COLS + n], va, vb[n], vl);
}
這裡用 _tu 是因為累加器要跨多輪 K-迴圈保留。當最後一輪 vl 變小時，尾端 lane 不能把舊累加器破壞掉。
decode 時，每個 query head 對當前 KV sequence 產生的 score 是長度為 seq_len 的向量，不是 prefill 那種 query-by-key 大矩陣。sequence length 變長時，kernel 仍不希望一次保留整條 score 向量，因此會在 KV block 上做分塊 softmax。
核心概念是維持兩個統計
m_prime: current max over scores
s_prime: exp(score - max) 的累積和
每個 KV block 算完 score 後， kernel 會更新 max 和 sum，並把前面累積的 v_prime 重新縮放。
接著做
v_prime = exp(score - m_new) @ V_block + old_v_prime * scale
這裡的 s_delta 是以新 running max 穩定化後的未正規化權重，不是已經除以 softmax denominator 的機率。舊累積值的 scale 是 exp(m_old - m_new)。每個非空 split 結束時會先用 running sum 正規化 partial value 並保存 LSE；decode_accumulate_kv_splits 再依各 split 的 LSE 加權合併，得到完整 sequence 的輸出。
程式碼裡會呼叫 index_gemm_kernel_nn
index_gemm_kernel_nn<scalar_t, kv_t, index_t>(
/* A */ s_delta,
/* B */ v_buffer + head_kv_id * v_strideH,
/* C */ v_prime,
/* ind */ cur_indices,
...
);
最後每個 KV split 會產生一個 partial output。decode_accumulate_kv_splits 再把 splits 合併成最後 output。
這就是 attn_logits 為什麼形狀是
[num_seqs, num_heads, num_kv_splits, head_size_v + 1]
最後多出來的 +1 用來放每個 split 的 log-sum-exp 相關值，合併 splits 時會用到。
INT8 decode 路徑會多傳
k_scale_buf
v_scale_buf
k_scale
v_scale
這裡有兩種不同的 scale 模式，需要分開讀
k_scale / v_scale 是靜態全局 scale，同一個值套用到對應 K 或 V 量化路徑。k_scale_buf / v_scale_buf 儲存動態 per-token / per-KV-head scale，kernel 會依 token slot 間接索引。新的 K/V 如果是 floating-point 資料型別，kernel 可依路徑計算動態 scale、quantize 成 int8，並寫入 scale buffer；如果輸入已經是 int8，則必須同時提供與該資料編碼相符的 scale。靜態全局 scale 與動態 buffer 不應被寫成同一種 per-token 機制。
在 dynamic per-token scale buffer 存在時，kernel 會在讀 KV block 前 pre-gather 每個 token 的縮放係數
k_scales_block[j] =
load_valid_int8_kv_scale(k_scale_buf, cur_indices[j], num_heads_kv, head_kv_id, "K");
這表示 INT8 attention 的效能不只看 int8 載入，縮放係數 buffer 的讀取、dequant、per-token 縮放係數資料配置都會影響 kernel 。
extend.cpp 的分塊設計Extend kernel 在 extend.cpp，一開始就定義了和 VLEN 相關的 block
static constexpr int EXTEND_BLOCK_N =
static_cast<int>(__riscv_v_fixed_vlen / 8);
static constexpr int BLOCK_M = 1024 / EXTEND_BLOCK_N;
這裡還需要實驗調參數目前先設 4KB。
Score tile s_i[BLOCK_M x BLOCK_N] 保持約 4KB
如果 K1 的 VLEN 是 256
EXTEND_BLOCK_N = 256 / 8 = 32
BLOCK_M = 1024 / 32 = 32
s_i = 32 x 32 x 4 bytes = 4096 bytes
目前只把 score tile s_i 控制在 4KB 左右，不是每個 thread 的全部 arena 只有 4KB。v_prime、Btmp、k_trans 和 v_buf 還會依 head size 額外佔用空間，這裡展示 VLEN 如何影響分塊，但總 scratch 大小要用 compute_buffer_size_per_thread 的完整配置來算。
compute_buffer_size_per_thread 會配置幾塊 buffer
s_i: float[BLOCK_M * BLOCK_N]
v_prime: float[BLOCK_M * head_size_v]
Btmp: scalar_t[BLOCK_N * max(head_size, head_size_v)]
k_trans: scalar_t[head_size * BLOCK_N]
v_buf: scalar_t[BLOCK_N * head_size_v]
這些 buffer 都在每個 thread 自己的 arena 裡，原因是 extend attention 會做 tiled QK、softmax、SV，過程中需要 temporary 分塊。
k_trans 這段值得看。KV cache 裡的 K資料配置通常不符合 kernel 最想要的存取形狀，extend kernel 會先把 K gather / transpose 成比較適合 GEMM 的 分塊，再做 QK。
Extend kernel 裡會先算幾個長度
seq_len = seq_lens[b];
extend_len = extend_seq_lens[b];
prefix_len = seq_len - extend_len;
Prefix 是已經在 KV cache 裡的 token。Current extend 是這一輪新進來的 token。
Stage 1 會處理 prefix
for n_start in prefix:
gather K from k_buffer through req_to_token
Q @ K
online softmax update
gather V from v_buffer
accumulate output
Stage 2 會處理 current extend token
for n_start in current extend:
read K from k_extend
apply causal mask
online softmax update
read V from v_extend
accumulate output
這裡和 decode 最大差別是：extend 一次處理多個 query row，所以它需要 BLOCK_M x BLOCK_N 的 score 分塊。
extend.cpp 裡的 softmax_update_row 是一個很好的 RVV 例子。
它先找 block 裡的 max
vfloat32m8_t v_max = __riscv_vfmv_v_f_f32m8(-INFINITY, vl_max);
...
v_max = __riscv_vfmax_vv_f32m8_tu(v_max, v_max, v_s, vl);
接著更新 global max
float m_new = std::max(m_acc, m_block);
float alpha = expf(m_acc - m_new);
m_acc = m_new;
然後把已經累積的 v_prime 乘上 alpha
v_vp = __riscv_vfmul_vf_f32m8(v_vp, alpha, vl);
最後把這一 block 的 score 轉成 exp，更新 l_acc
v_s = __riscv_vfsub_vf_f32m8(v_s, m_new, vl);
vfloat32m8_t v_exp = vfexp_f32m8(v_s, vl);
...
l_acc = l_acc * alpha + l_block;
這段把昨天的 reduction、近似函式、尾端 handling 都放進 attention，softmax 的數值穩定性靠 m_acc 和 l_acc 維持，各個 block 的結果還要合併到同一組 running max / running sum。
Attention效能測試需要比 linear 更小心，因為它同時受多個因素影響
req_to_token 的 access 模式看效能分析器時，可以先確認
sgl-kernel::decode_attention_cpu
sgl-kernel::extend_attention_cpu
sgl-kernel::decode_attention_int8_cpu
sgl-kernel::extend_attention_int8_cpu
然後再看單批次 prefill / decode 延遲。Decode 和 extend 的 bottleneck 可能不同。Decode 常受 KV cache 串流讀取和小批次額外成本 影響，extend 則更看 分塊 buffer、QK/SV GEMM、softmax update
今天把 SGLang RVV attention kernel 的主線走完
BLOCK_M x BLOCK_N 分塊req_to_token、req_pool_indices、seq_lens 把 SGLang 請求狀態交給 kernelEXTEND_BLOCK_N = VLEN / 8，K1 VLEN=256 時 block_N 會是 32，score 分塊約 4KBm_acc / l_acc，並在每個 block 重新縮放已累積的 value到這裡我們把 SGLang RVV 後端的整合和主要 AI 運算子都已經看過，發現手寫 RVV 運算子還蠻累人的，所以明天開始進入到 AI compiler 的範圍，看 torch.compile 和 PyTorch Inductor，看看能不能讓編譯器自動產生 RVV C++，減少手寫自訂運算子的負擔。
python/sglang/srt/layers/attention/rvv_backend.py
sgl-kernel/csrc/cpu/riscv64/decode.cpp
sgl-kernel/csrc/cpu/riscv64/extend.cpp
iThome鐵人賽
看影片追技術